/* Find all lattice points on a sphere with a given radius, and compute
 * the sum of their Manhattan distance to the origin.
 */

/*
Note that for type III points, the number of points is equal to 
\lfloor r / 10 \rfloor
This surprising result should be used to improve the algorithm.

Type I  : 10000000000
Type II : 255994421248
Checking primitive lattice for r = 1...
Found 0 points.
Checking primitive lattice for r = 5...
Found 0 points.
Checking primitive lattice for r = 25...
Found 2 points.
Checking primitive lattice for r = 125...
Found 12 points.
Checking primitive lattice for r = 625...
Found 62 points.
Checking primitive lattice for r = 3125...
Found 312 points.
Checking primitive lattice for r = 15625...
Found 1562 points.
Checking primitive lattice for r = 78125...
Found 7812 points.
Checking primitive lattice for r = 390625...
Found 39062 points.
Checking primitive lattice for r = 1953125...
Found 195312 points.
Checking primitive lattice for r = 9765625...
Found 976562 points.
Type III: 109852810307776512
*/

#include <iostream>
#include <algorithm>
#include "euler/divisor.hpp"
#include "euler/lattice.hpp"
#include "euler.h"

BEGIN_PROBLEM(360, solve_problem_360)
	PROBLEM_TITLE("Scary Sphere")
	PROBLEM_ANSWER("878825614395267072")
	PROBLEM_DIFFICULTY(2)
	PROBLEM_FUN_LEVEL(2)
	PROBLEM_TIME_COMPLEXITY("r*ln(r)")
	PROBLEM_SPACE_COMPLEXITY("r")
	PROBLEM_KEYWORDS("lattice point,sphere")
END_PROBLEM()

// Method 0: Brute force. This method requires no restriction on R.
#if 0
static long long S0(int r)
{
	long long l = 0;
	for (int a = -r; a <= r; a++)
	{
		for (int b = -r; b <= r; b++)
		{
			for (int c = -r; c <= r; c++)
			{
				if (a*a + b*b + c*c == r*r)
				{
					// std::cout << a << ", " << b << ", " << c << std::endl;
					l += abs(a) + abs(b) + abs(c);
				}
			}
		}
	}
	return l;
}
#else
static long long S0(int r)
{
	long long L1 = r;

	long long L2 = 0;
	for (int a = 1; a < r; a++)
	{
		for (int b = 1; b < r; b++)
		{
			if (a*a + b*b == r*r)
			{
				L2 += (a + b);
			}
		}
	}

	long long L3 = 0;
	for (int a = 1; a < r; a++)
	{
		for (int b = 1; b < r; b++)
		{
			for (int c = 1; c < r; c++)
			{
				if (a*a + b*b + c*c == r*r)
				{
					L3 += (a + b + c);
				}
			}
		}
	}

	std::cout << "Type I  : " << L1 << std::endl;
	std::cout << "Type II : " << L2 << std::endl;
	std::cout << "Type III: " << L3 << std::endl;
	return 6*L1 + 12*L2 + 8*L3;
}
#endif

static void sort3(int &a, int &b, int &c)
{
	if (a > b) std::swap(a, b);
	if (b > c) std::swap(b, c);
	if (a > b) std::swap(a, b);
}

// Method 1: Enumerate Pythagorean quadruple. This method imposes no 
// restriction on R.
static long long S1(long long r, bool verbose)
{
	// Split r into even and odd factors.
	int r_even = 1;
	while (r % 2 == 0)
	{
		r_even *= 2;
		r /= 2;
	}
	int r_odd = (int)r;

	// Prepare counters.
	int c1 = 0, c2 = 0, c3 = 0;
	long long L1 = 0, L2 = 0, L3 = 0;

	// Enumerates primitive lattice points on each sphere with odd radius.
	std::for_each(euler::divisor_iterator<int>(r_odd), euler::divisor_iterator<int>(), 
		[&](int d) {
			if (verbose)
				std::cout << "Checking primitive lattice for r = " << d << "..." << std::endl;
			long long k = (r_odd / d);
			euler::sphere_primitive_lattice(d, [&](int a, int b, int c) {
				sort3(a, b, c);
				if (a == 0 && b == 0)
				{
					c1++;
					L1 += k * c;
				}
				else if (a == 0)
				{
					c2++;
					L2 += k * 2 * (b + c);
				}
				else if (a != b && b != c && c != a)
				{
					c3++;
					L3 += k * 6 * (a + b + c);
				}
				else
				{
					c3++;
					L3 += k * 3 * (a + b + c);
				}
			});
			//if (verbose)
			//	std::cout << "Found " << count << " points." << std::endl;
	});
	if (verbose)
	{
		std::cout << "Type I  : " << L1*r_even << " (" << c1 << " points)" << std::endl;
		std::cout << "Type II : " << L2*r_even << " (" << c2 << " points)" << std::endl;
		std::cout << "Type III: " << L3*r_even << " (" << c3 << " points)" << std::endl;
	}
	return (6*L1 + 12*L2 + 8*L3) * r_even;
}

// Method 2: Make use of the sum of squares function, and note that 
// only summing the z-value is needed to compute the Manhattan distance.
// This method imposes no restriction on R.
// NOT IMPLEMENTED

// Method 3: Make use of the fact that the radius is (2*5)^10 to generate
// the lattice points recursively. This only works for the radius of the
// form (2*5)^n.
// NOT IMPLEMENTED

static void solve_problem_360()
{
#if 0
	//const int r = 45;
	const int r = 100;
	//const int r = 100000000;
#else
	const long long r = 10000000000;
#endif

	bool verbose = false;

	//std::cout << S0(r) << std::endl;
	std::cout << S1(r, verbose) << std::endl;
}
